# -*- coding: utf-8 -*-
###################################################
#    This file is part of blockIMH.
#
#    blockIMH is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    blockIMH is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with blockIMH.  If not, see <http://www.gnu.org/licenses/>.
###################################################
#! /usr/bin/env python

from __future__ import division
import os
import sys
import scipy.weave as weave
from scipy.stats import norm, cauchy, uniform
from scipy import special
from numpy import random, exp, sqrt, power, log, newaxis, zeros, \
        zeros_like, float32, int32, repeat, minimum, ones, mean, average, sum, \
        min, prod, nan, load, vstack, array
try:
    import pycuda
    import pycuda.autoinit
    import pycuda.gpuarray as gpuarray
    import pycuda.driver as cuda
    import pycuda.cumath as cumath
    from pycuda.compiler import SourceModule
    from mycudafunctions import *
except:
    CUDA_ENABLED = False
import permutations
import raoblackwell
from probit import rproposal, fastomega


def IMHsteps(p, omegaproposalsblock, omegacurrent, uniformsblock, weightsRB1, nboccurs, chain, proposalsblock):
    parameterdimension = chain.shape[2]
    code = \
    """
    float acceptratio;
    float logacceptratio;
    int currentidx;
    int acceptation;
    float wcurrent;
    for(int k = 0; k < p; k++)
    {
        currentidx = 0;
        wcurrent = omegacurrent;
        for(int t = 0; t < p; t++)
        {
            logacceptratio = (double) omegaproposalsblock(t, k) - (double) wcurrent;
            if (logacceptratio > 0){ logacceptratio = 0;}
            acceptation = (log(uniformsblock(t, k)) < logacceptratio);
            acceptratio = exp(logacceptratio);
            weightsRB1(currentidx, k) += 1 - acceptratio;
            weightsRB1(t + 1, k) += acceptratio;
            if (acceptation){
                //acceptcounter(0) += 1;
                nboccurs(t + 1, k) += 1;
                currentidx = t + 1 ;
                for (int param = 0; param < parameterdimension; param ++){
                    chain(t + 1, k, param) = proposalsblock(t, k, param);
                }
                wcurrent = omegaproposalsblock(t, k);
            }
            else{
                //rejectcounter(0) += 1;
                nboccurs(currentidx, k) += 1;
                for (int param = 0; param < parameterdimension; param ++){
                    chain(t + 1, k, param) = chain(t, k, param);
                }
            }
        }
    }
    """
    weave.inline(code,["p", "omegaproposalsblock", "omegacurrent", "uniformsblock", "weightsRB1", \
            "nboccurs", "chain", "proposalsblock", "parameterdimension"], type_converters=weave.converters.blitz)

def parallelchains(parameterdimension, h, p, proposals, permutations, uniformsblock, omegaproposals, x_starting, omegacurrent, RB2 = False):
    chain = zeros((p+1, p, parameterdimension))
    chain[0, ...] = x_starting
    proposalsblock = proposals[permutations]
    omegaproposalsblock = omegaproposals[permutations]
    diffvalues = zeros((p+1, p, parameterdimension))
    for i in range(parameterdimension):
        diffvalues[0, :, i] = repeat(x_starting[i], p)
    diffvalues[1:(p+1), ...] = proposalsblock
    weightsRB1 = zeros((p+1, p))
    nboccurs = zeros((p+1, p))
    IMHsteps(p, omegaproposalsblock, omegacurrent, uniformsblock, weightsRB1, nboccurs, chain, proposalsblock)
    blockestim = zeros(parameterdimension)
    estimRB1 = zeros(parameterdimension)
    for i in range(parameterdimension):
        blockestim[i] = average(h(diffvalues[..., i]), weights = nboccurs)
        estimRB1[i] = average(h(diffvalues[..., i]), weights = weightsRB1)
    if RB2:
        estimRB2 = zeros(parameterdimension)
        weightsRB2 = zeros((p+1, p))
        expomega = exp(omegaproposalsblock)
        expomegacurrent = exp(omegacurrent)
        for k in range(p):
            raoblackwell.raoblackwell(k, p, expomega, expomegacurrent, weightsRB2)
        for i in range(parameterdimension):
            estimRB2[i] = average(h(diffvalues[..., i]), weights = weightsRB2)
    else:
        estimRB2 = zeros(parameterdimension)
    return({"blockestim": blockestim, "estimRB1": estimRB1, "acceptedvalues": chain[1:(p+1), ...], "estimRB2" : estimRB2})

def blockIMH(x_starting, T, p, h):
    parameterdimension = x_starting.shape[0]
    x = zeros((T + 1, parameterdimension))
    x[0, :] = x_starting
    omegacurrent = fastomega(x_starting[newaxis, :])[0]
    nbblocks = T // p
    acceptedvalues = zeros((p, p, parameterdimension))
    uniformsblock = zeros((p,p))
    blockestimRandPermut = zeros((nbblocks, parameterdimension))
    estimRB1RandPermut = zeros((nbblocks, parameterdimension))
    estimRB2RandPermut = zeros((nbblocks, parameterdimension))
    allproposals = rproposal(size = p * nbblocks)
    allomegaproposals = fastomega(allproposals)
    estimIS = average(h(allproposals), weights = exp(allomegaproposals), axis = 0)
    for block in range(nbblocks):
        permutationmatrix = permutations.createRandomPermutation(p)
        #print "block #%i" % block
        y = allproposals[(block * p):((block + 1) * p), :]
        omegaproposals = allomegaproposals[(block * p):((block + 1) * p)]
        uniformsblock[...] = random.uniform(size = p**2).reshape(p,p)
### Random order
        resultsRandPermut = parallelchains(parameterdimension, h, p, y, permutationmatrix, uniformsblock, \
                omegaproposals, x_starting, omegacurrent, RB2 = True)
        blockestimRandPermut[block, :] = resultsRandPermut["blockestim"]
        estimRB1RandPermut[block, :] = resultsRandPermut["estimRB1"]
        estimRB2RandPermut[block, :] = resultsRandPermut["estimRB2"]
## Prepare the next block
        selected = int32(random.uniform(size = 1) * p)
        x[(p * block + 1):(p * (block+1) +1), :] = resultsRandPermut["acceptedvalues"][:, selected, :][:, 0, :]
        x_starting = x[p * (block + 1), :]
        omegacurrent = fastomega(x_starting[newaxis, :])[0]
    chain = x[1:(T+1), :]
    return {"chain": chain, "basicestim": mean(h(chain), axis = 0), \
            "blockestimRandPermut": mean(blockestimRandPermut, axis = 0), \
            "estimRB1RandPermut": mean(estimRB1RandPermut, axis = 0), \
            "estimRB2RandPermut": mean(estimRB2RandPermut, axis = 0), \
            "estimIS": estimIS}

